package xxx;


import com.lxw.mutils.FileUtils;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameUtils;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.Set;

public class VideoTextExtractor {
    // Tesseract-OCR安装路径
    public static final String pathToTessdataFolder = "D:\\Program Files\\Tesseract-OCR\\tessdata\\";
    // 加载视频
    public static final String pathToVideoFile = "E:\\图灵学院\\图灵四期\\06 微信亿级IM后台系统\\01 企业级IM系统后端架构实战(上)-诸葛.mp4";
    // 解析结果
    public static final String resultFile = "E:\\tmp\\tmp.txt";


    public static void main(String[] args) throws TesseractException {
        extracted();
    }


    private static void extracted() {
        // 设置Tesseract OCR库的路径
        File tessDataFolder = new File(pathToTessdataFolder);
        System.setProperty("TESSDATA_PREFIX", tessDataFolder.getAbsolutePath());
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(pathToVideoFile);
        try {
            grabber.start();
            Set<String> set = new LinkedHashSet<>();
            // 遍历视频帧
            int lengthInFrames = grabber.getLengthInFrames();
            for (int i = 0; i < lengthInFrames; i++) {
                System.out.println("进度 " + i + " / " + lengthInFrames);
                try {
                    Frame frame = grabber.grabImage();
                    BufferedImage bufferedImage = Java2DFrameUtils.toBufferedImage(frame);

                    // 将帧转换为灰度图像
                    BufferedImage grayImage = new BufferedImage(bufferedImage.getWidth(), bufferedImage.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
                    Graphics2D graphics = grayImage.createGraphics();
                    graphics.drawImage(bufferedImage, 0, 0, null);
                    graphics.dispose();

                    // 创建临时文件保存图像
                    File tempImageFile = File.createTempFile("frame", ".png");
                    ImageIO.write(grayImage, "png", tempImageFile);

                    Tesseract tesseract = getTesseract(tessDataFolder);
                    String result = tesseract.doOCR(tempImageFile);
                    set.add(result);
                    // 删除临时文件
                    tempImageFile.delete();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            File file = new File(resultFile);

            FileUtils.write2File(file, new ArrayList<>(set));

            grabber.stop();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static Tesseract getTesseract(File tessDataFolder) {
        // 使用Tesseract OCR进行文字识别
        Tesseract tesseract = new Tesseract();
        // 设置中文训练库 中文准确率还是不如百度Orc
        tesseract.setLanguage("chi_sim");
        tesseract.setDatapath(tessDataFolder.getAbsolutePath());
        return tesseract;
    }
}
